import time
import sys

from unitree_sdk2py.core.channel import ChannelPublisher, ChannelFactoryInitialize
from unitree_sdk2py.core.channel import ChannelSubscriber, ChannelFactoryInitialize
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowCmd_
from unitree_sdk2py.idl.default import unitree_hg_msg_dds__LowState_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowCmd_
from unitree_sdk2py.idl.unitree_hg.msg.dds_ import LowState_
from unitree_sdk2py.utils.crc import CRC
from unitree_sdk2py.utils.thread import RecurrentThread
from unitree_sdk2py.comm.motion_switcher.motion_switcher_client import MotionSwitcherClient

import numpy as np

kPi = 3.141592654
kPi_2 = 1.57079632

class G1JointIndex:
    # Left leg
    LeftHipPitch = 0
    LeftHipRoll = 1
    LeftHipYaw = 2
    LeftKnee = 3
    LeftAnklePitch = 4
    LeftAnkleB = 4
    LeftAnkleRoll = 5
    LeftAnkleA = 5

    # Right leg
    RightHipPitch = 6
    RightHipRoll = 7
    RightHipYaw = 8
    RightKnee = 9
    RightAnklePitch = 10
    RightAnkleB = 10
    RightAnkleRoll = 11
    RightAnkleA = 11

    WaistYaw = 12
    WaistRoll = 13        # NOTE: INVALID for g1 23dof/29dof with waist locked
    WaistA = 13           # NOTE: INVALID for g1 23dof/29dof with waist locked
    WaistPitch = 14       # NOTE: INVALID for g1 23dof/29dof with waist locked
    WaistB = 14           # NOTE: INVALID for g1 23dof/29dof with waist locked

    # Left arm
    LeftShoulderPitch = 15
    LeftShoulderRoll = 16
    LeftShoulderYaw = 17
    LeftElbow = 18
    LeftWristRoll = 19
    LeftWristPitch = 20   # NOTE: INVALID for g1 23dof
    LeftWristYaw = 21     # NOTE: INVALID for g1 23dof

    # Right arm
    RightShoulderPitch = 22
    RightShoulderRoll = 23
    RightShoulderYaw = 24
    RightElbow = 25
    RightWristRoll = 26
    RightWristPitch = 27  # NOTE: INVALID for g1 23dof
    RightWristYaw = 28    # NOTE: INVALID for g1 23dof

    kNotUsedJoint = 29 # NOTE: Weight

class Custom:
    def __init__(self):
        self.time_ = 0.0
        self.control_dt_ = 0.02  
        self.duration_ = 3.0   
        self.counter_ = 0
        self.weight = 0.
        self.weight_rate = 0.2
        self.kp = 60.
        self.kd = 1.5
        self.dq = 0.
        self.tau_ff = 0.
        self.mode_machine_ = 0
        self.low_cmd = unitree_hg_msg_dds__LowCmd_()  
        self.low_state = None 
        self.first_update_low_state = False
        self.crc = CRC()
        self.done = False

        self.target_pos = [
            0., kPi_2,  0., kPi_2, 0., 0., 0.,
            0., -kPi_2, 0., kPi_2, 0., 0., 0., 
            0, 0, 0
        ]

        self.arm_joints = [
          G1JointIndex.LeftShoulderPitch,  G1JointIndex.LeftShoulderRoll,
          G1JointIndex.LeftShoulderYaw,    G1JointIndex.LeftElbow,
          G1JointIndex.LeftWristRoll,      G1JointIndex.LeftWristPitch,
          G1JointIndex.LeftWristYaw,
          G1JointIndex.RightShoulderPitch, G1JointIndex.RightShoulderRoll,
          G1JointIndex.RightShoulderYaw,   G1JointIndex.RightElbow,
          G1JointIndex.RightWristRoll,     G1JointIndex.RightWristPitch,
          G1JointIndex.RightWristYaw,
          G1JointIndex.WaistYaw,
          G1JointIndex.WaistRoll,
          G1JointIndex.WaistPitch
        ]

    def Init(self):
        # create publisher #
        self.arm_sdk_publisher = ChannelPublisher("rt/arm_sdk", LowCmd_)
        self.arm_sdk_publisher.Init()

        # create subscriber # 
        self.lowstate_subscriber = ChannelSubscriber("rt/lowstate", LowState_)
        self.lowstate_subscriber.Init(self.LowStateHandler, 10)

    def Start(self):
        self.lowCmdWriteThreadPtr = RecurrentThread(
            interval=self.control_dt_, target=self.LowCmdWrite, name="control"
        )
        while self.first_update_low_state == False:
            time.sleep(1)

        if self.first_update_low_state == True:
            self.lowCmdWriteThreadPtr.Start()

    def LowStateHandler(self, msg: LowState_):
        self.low_state = msg

        if self.first_update_low_state == False:
            self.first_update_low_state = True
        
    def LowCmdWrite(self):
        self.time_ += self.control_dt_

        if self.time_ < self.duration_ :
          # [Stage 1]: set robot to zero posture
          self.low_cmd.motor_cmd[G1JointIndex.kNotUsedJoint].q =  1 # 1:Enable arm_sdk, 0:Disable arm_sdk
          for i,joint in enumerate(self.arm_joints):
            ratio = np.clip(self.time_ / self.duration_, 0.0, 1.0)
            self.low_cmd.motor_cmd[joint].tau = 0. 
            self.low_cmd.motor_cmd[joint].q = (1.0 - ratio) * self.low_state.motor_state[joint].q 
            self.low_cmd.motor_cmd[joint].dq = 0. 
            self.low_cmd.motor_cmd[joint].kp = self.kp 
            self.low_cmd.motor_cmd[joint].kd = self.kd

        elif self.time_ < self.duration_ * 3 :
          # [Stage 2]: lift arms up
          for i,joint in enumerate(self.arm_joints):
              ratio = np.clip((self.time_ - self.duration_) / (self.duration_ * 2), 0.0, 1.0)
              self.low_cmd.motor_cmd[joint].tau = 0. 
              self.low_cmd.motor_cmd[joint].q = ratio * self.target_pos[i] + (1.0 - ratio) * self.low_state.motor_state[joint].q 
              self.low_cmd.motor_cmd[joint].dq = 0. 
              self.low_cmd.motor_cmd[joint].kp = self.kp 
              self.low_cmd.motor_cmd[joint].kd = self.kd

        elif self.time_ < self.duration_ * 6 :
          # [Stage 3]: set robot back to zero posture
          for i,joint in enumerate(self.arm_joints):
              ratio = np.clip((self.time_ - self.duration_*3) / (self.duration_ * 3), 0.0, 1.0)
              self.low_cmd.motor_cmd[joint].tau = 0. 
              self.low_cmd.motor_cmd[joint].q = (1.0 - ratio) * self.low_state.motor_state[joint].q
              self.low_cmd.motor_cmd[joint].dq = 0. 
              self.low_cmd.motor_cmd[joint].kp = self.kp 
              self.low_cmd.motor_cmd[joint].kd = self.kd

        elif self.time_ < self.duration_ * 7 :
          # [Stage 4]: release arm_sdk
          for i,joint in enumerate(self.arm_joints):
              ratio = np.clip((self.time_ - self.duration_*6) / (self.duration_), 0.0, 1.0)
              self.low_cmd.motor_cmd[G1JointIndex.kNotUsedJoint].q =  (1 - ratio) # 1:Enable arm_sdk, 0:Disable arm_sdk
        
        else:
            self.low_cmd.motor_cmd[joint].kp = self.kp
            self.low_cmd.motor_cmd[joint].kd = self.kdself.done = True
  
        self.low_cmd.crc = self.crc.Crc(self.low_cmd)
        self.arm_sdk_publisher.Write(self.low_cmd)

class CustomNew(Custom):
    def __init__(self):
        super().__init__()
        self.target_pos = {
            # 左手目标坐标
            self.arm_joints[0]: 0.0,        # LeftShoulderPitch
            self.arm_joints[1]: 0.0,        # LeftShoulderRoll
            self.arm_joints[2]: 0.0,        # LeftShoulderYaw
            self.arm_joints[3]: kPi_2*2/3,  # LeftElbow
            self.arm_joints[4]: 0.0,        # LeftWristRoll
            self.arm_joints[5]: 0.0,        # LeftWristPitch
            self.arm_joints[6]: 0.0,        # LeftWristYaw
            # 右手目标坐标
            self.arm_joints[7]: 0.0,        # RightShoulderPitch
            self.arm_joints[8]: 0.0,        # RightShoulderRoll
            self.arm_joints[9]: 0.0,        # RightShoulderYaw
            self.arm_joints[10]: 0.0,       # RightElbow
            self.arm_joints[11]: -kPi_2,    # RightWristRoll
            self.arm_joints[12]: 0.0,       # RightWristPitch
            self.arm_joints[13]: 0.0,       # RightWristYaw
            # 腰部目标坐标
            self.arm_joints[14]: 0.0,       # WaistYaw
            self.arm_joints[15]: 0.0,       # WaistRoll
            self.arm_joints[16]: 0.0,       # WaistPitch
        }


    def LowCmdWrite(self):
        self.time_ += self.control_dt_

        if self.time_ < self.duration_:
            # [阶段1]: 旋转手掌向下，手腕下翻
            self.low_cmd.motor_cmd[G1JointIndex.kNotUsedJoint].q = 1  # 1:启用 arm_sdk, 0:禁用 arm_sdk
            for joint, target in self.target_pos.items():
                ratio = np.clip(self.time_ / self.duration_, 0.0, 1.0)
                self.low_cmd.motor_cmd[joint].tau = 0.
                self.low_cmd.motor_cmd[joint].q = ratio * target + (1.0 - ratio) * self.low_state.motor_state[joint].q  # 插值计算目标位置
                self.low_cmd.motor_cmd[joint].dq = 0.
                self.low_cmd.motor_cmd[joint].kp = self.kp
                self.low_cmd.motor_cmd[joint].kd = self.kd

        # elif self.time_ < self.duration_ * 3:
        #     # [阶段2]: 旋转手掌向下，手腕下翻
        #     for joint, target in self.target_pos.items():
        #         ratio = np.clip((self.time_ - self.duration_) / (self.duration_ * 2), 0.0, 1.0)
        #         self.low_cmd.motor_cmd[joint].tau = 0.
        #         self.low_cmd.motor_cmd[joint].q = ratio * target + (1.0 - ratio) * self.low_state.motor_state[joint].q  # 插值计算目标位置
        #         self.low_cmd.motor_cmd[joint].dq = 0.
        #         self.low_cmd.motor_cmd[joint].kp = self.kp
        #         self.low_cmd.motor_cmd[joint].kd = self.kd

        # elif self.time_ < self.duration_ * 6:
        #     ratio = np.clip((self.time_ - self.duration_ * 3) / (self.duration_ * 3), 0.0, 1.0)
        #     self.low_cmd.motor_cmd[G1JointIndex.RightShoulderYaw].tau = 0.
        #     self.low_cmd.motor_cmd[G1JointIndex.RightShoulderYaw].q = ratio * -kPi_2/2 + (1.0 - ratio) * self.low_state.motor_state[G1JointIndex.RightShoulderYaw].q  # 插值计算目标位置
        #     self.low_cmd.motor_cmd[G1JointIndex.RightShoulderYaw].dq = 0.
        #     self.low_cmd.motor_cmd[G1JointIndex.RightShoulderYaw].kp = self.kp
        #     self.low_cmd.motor_cmd[G1JointIndex.RightShoulderYaw].kd = self.kd
        
        # elif self.time_ < self.duration_ * 6:
        #     ratio = np.clip((self.time_ - self.duration_ * 3) / (self.duration_ * 3), 0.0, 1.0)
        #     self.low_cmd.motor_cmd[G1JointIndex.LeftElbow].tau = 0.
        #     self.low_cmd.motor_cmd[G1JointIndex.LeftElbow].q = ratio * kPi_2/2 + (1.0 - ratio) * self.low_state.motor_state[G1JointIndex.LeftElbow].q  # 插值计算目标位置
        #     self.low_cmd.motor_cmd[G1JointIndex.LeftElbow].dq = 0.
        #     self.low_cmd.motor_cmd[G1JointIndex.LeftElbow].kp = self.kp
        #     self.low_cmd.motor_cmd[G1JointIndex.LeftElbow].kd = self.kd


        # elif self.time_ < self.duration_ * 3:
        #     # [阶段2]: 抬起手臂
        #     for i, joint in enumerate(self.arm_joints):
        #         ratio = np.clip((self.time_ - self.duration_) / (self.duration_ * 2), 0.0, 1.0)
        #         self.self.low_cmd.motor_cmd[joint].kp = self.kp
                # self.low_cmd.motor_cmd[joint].kd = self.kdlow_cmd.motor_cmd[joint].kp = self.kp  # 设置位置增益
        #         self.low_cmd.motor_cmd[joint].kd = self.kd  # 设置速度增益

        # elif self.time_ < self.duration_ * 6:
        #     # [阶段3]: 将机器人恢复到零姿态
        #     for i, joint in enumerate(self.arm_joints):
        #         ratio = np.clip((self.time_ - self.duration_ * 3) / (self.duration_ * 3), 0.0, 1.0)
        #         self.low_cmd.motor_cmd[joint].tau = 0.  # 力矩设为0
        #         self.low_cmd.motor_cmd[joint].q = (1.0 - ratio) * self.low_state.motor_state[joint].q  # 插值计算目标位置
        #         self.low_cmd.motor_cmd[joint].dq = 0.  # 速度设为0
        #         self.low_cmd.motor_cmd[joint].kp = self.kp  # 设置位置增益
        #         self.low_cmd.motor_cmd[joint].kd = self.kd  # 设置速度增益

        elif self.time_ < self.duration_ * 7:
            # [阶段4]: 释放 arm_sdk
            for i, joint in enumerate(self.arm_joints):
                ratio = np.clip((self.time_ - self.duration_ * 6) / self.duration_, 0.0, 1.0)
                self.low_cmd.motor_cmd[G1JointIndex.kNotUsedJoint].q = (1 - ratio)  # 1:启用 arm_sdk, 0:禁用 arm_sdk

        else:
            # 动作完成
            self.done = True

        # 计算 CRC 校验码并写入命令
        self.low_cmd.crc = self.crc.Crc(self.low_cmd)
        self.arm_sdk_publisher.Write(self.low_cmd)


    def MoveJoints(self, duration, joint_targets):
        """
        控制指定关节在某个时间内完成动作。

        :param duration: 完成动作的时间（秒）。
        :param joint_targets: 字典，键为关节索引，值为目标运动幅度（弧度）。
        """
        start_time = self.time_  # 记录开始时间
        end_time = start_time + duration  # 计算结束时间
        start_state = { }
        for i,joint in enumerate(self.arm_joints):
            start_state.setdefault(joint, self.low_state.motor_state[joint].q)

        while self.time_ < end_time:
            self.low_cmd.motor_cmd[G1JointIndex.kNotUsedJoint].q = 1 
            ratio = np.clip((self.time_ - start_time) / duration, 0.0, 1.0)  # 计算插值比例
            for joint in self.arm_joints:
                if joint in joint_targets:
                    # 插值计算目标位置
                    target = joint_targets[joint]
                    self.low_cmd.motor_cmd[joint].tau = 0.0  # 力矩为0
                    self.low_cmd.motor_cmd[joint].q = ratio * target + (1.0 - ratio) * self.low_state.motor_state[joint].q
                else:
                    # 保持当前状态
                    self.low_cmd.motor_cmd[joint].q = start_state[joint]
                    self.low_cmd.motor_cmd[joint].tau = 0.0  # 力矩为0
                self.low_cmd.motor_cmd[joint].dq = 0.0  # 速度为0
                self.low_cmd.motor_cmd[joint].kp = self.kp  # 设置位置增益
                self.low_cmd.motor_cmd[joint].kd = self.kd  # 设置速度增益

            # 计算 CRC 校验码并发送命令
            self.low_cmd.crc = self.crc.Crc(self.low_cmd)
            self.arm_sdk_publisher.Write(self.low_cmd)

            time.sleep(self.control_dt_)  # 等待下一个控制周期
            self.time_ += self.control_dt_  # 更新时间

        print(f"动作完成，关节目标已达到：{joint_targets}")
    
    def StartMoveJoints(self, duration, joint_targets):
        """
        调用 MoveJoints 方法以控制指定关节完成动作。

        :param duration: 完成动作的时间（秒）。
        :param joint_targets: 字典，键为关节索引，值为目标运动幅度（弧度）。
        """
        self.lowCmdWriteThreadPtr = RecurrentThread(
            interval=self.control_dt_, target=lambda: self.MoveJoints(duration, joint_targets), name="control"
        )
        while self.first_update_low_state == False:
            time.sleep(1)

        if self.first_update_low_state == True:
            self.lowCmdWriteThreadPtr.Start()


if __name__ == '__main__':

    print("WARNING: Please ensure there are no obstacles around the robot while running this example.")
    input("Press Enself.low_cmd.motor_cmdter to continue...")

    if len(sys.argv)>1:
        ChannelFactoryInitialize(0, sys.argv[1])
    else:
        ChannelFactoryInitialize(0)


    test_target = {
        G1JointIndex.RightWristRoll:kPi_2,
    }

    custom = CustomNew()
    custom.Init()
    # custom.Start()
    custom.StartMoveJoints(3, test_target)

    while True:        
        time.sleep(1)
        if custom.done: 
           print("Done!")
           break
        #    sys.exit(-1)    